package org.javaboy.custom_limit_demo.service;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.util.Set;
import java.util.UUID;

@Service
public class LimitService {

    @Autowired
    StringRedisTemplate redisTemplate;

    public static final String LIMIT_KEY = "javaboy_limit";

    /**
     * 限流方法，每个请求到达之后，先调用这个方法判断是否允许请求
     * <p>
     * 假设目前的限流规则是 3s 内可以处理 5 个请求
     *
     * @return
     */
    public boolean isAllowed() {
        //1. 获取当前时间戳
        long currentTimeMillis = System.currentTimeMillis();
        if (redisTemplate.hasKey(LIMIT_KEY)) {
            //方式一
            //2. 获取当前时间窗内的所有请求
//            Set<String> requests = redisTemplate.opsForZSet().rangeByScore(LIMIT_KEY, currentTimeMillis - 1000 * 3, currentTimeMillis);
            //方式二
            //首先，删除当前时间窗之前的请求
            redisTemplate.opsForZSet().removeRangeByScore(LIMIT_KEY, 0, currentTimeMillis - 1000 * 3);
            //2. 获取当前时间窗内的所有请求（也就是剩余的所有请求）
            Set<String> requests = redisTemplate.opsForZSet().rangeByScore(LIMIT_KEY, 0, currentTimeMillis);
            //3. 判断当前时间窗内的请求是否超过5个
            //如果超过5个，则不允许请求
            if (requests != null && requests.size() > 4) {
                return false;
            }
        }
        //最后，将当前请求存入到时间轴上
        //score 就是当前的时间戳
        redisTemplate.opsForZSet().add(LIMIT_KEY, UUID.randomUUID().toString(), currentTimeMillis);
        return true;
    }
}
